#!/root/anaconda3/envs/DCE2/bin/python
# -*- coding: UTF-8 -*-
import os
import rospy
import numpy as np
from sensor_msgs.msg import Imu
import matplotlib.pyplot as plt
from cv_bridge import CvBridge
import uuid
import threading

bridge = CvBridge()

def NORM(x):
    return np.math.sqrt(np.dot(x,x))

def dis_ACC(x, y):
    scale = 0.5
    diff = x - y
    dis = np.dot(x, y) / (NORM(x) * NORM(y) + 1e-8)
    dis = (1 - scale * dis) * NORM(diff)
    return dis


def estimate_dtw(A, B, dis_func=dis_ACC):
    N_A = len(A)
    N_B = len(B)

    D = np.zeros([N_A, N_B])
    D[0, 0] = dis_func(A[0], B[0])

    #左边一列
    for i in range(1, N_A):
        D[i, 0] = D[i - 1, 0] + dis_func(A[i], B[0])
    #下边一行
    for j in range(1, N_B):
        D[0, j] = D[0, j - 1] + dis_func(A[0], B[j])
    #中间部分
    for i in range(1, N_A):
        for j in range(1, N_B):
            D[i, j] = dis_func(A[i], B[j]) + min(D[i - 1, j], D[i, j - 1], D[i - 1, j - 1])

    # 路径回溯
    i = N_A - 1
    j = N_B - 1
    count = 0
    d = np.zeros(max(N_A, N_B) * 3)
    path = []
    while True:
        if i > 0 and j > 0:
            path.append((i, j))
            m = min(D[i - 1, j], D[i, j - 1], D[i - 1, j - 1])
            if m == D[i - 1, j - 1]:
                d[count] = D[i, j] - D[i - 1, j - 1]
                i = i - 1
                j = j - 1
                count = count + 1

            elif m == D[i, j - 1]:
                d[count] = D[i, j] - D[i, j - 1]
                j = j - 1
                count = count + 1

            elif m == D[i - 1, j]:
                d[count] = D[i, j] - D[i - 1, j]
                i = i - 1
                count = count + 1

        elif i == 0 and j == 0:
            path.append((i, j))
            d[count] = D[i, j]
            count = count + 1
            break

        elif i == 0:
            path.append((i, j))
            d[count] = D[i, j] - D[i, j - 1]
            j = j - 1
            count = count + 1

        elif j == 0:
            path.append((i, j))
            d[count] = D[i, j] - D[i - 1, j]
            i = i - 1
            count = count + 1

    mean = np.sum(d) / count
    return mean, path[::-1], D

datalist=[]

imu_buf=threading.Lock()

def handleIMU(data):
    imudata=[]
    imu_buf.acquire()
    imudata.append(data.linear_acceleration.x)
    imudata.append(data.linear_acceleration.y)
    imudata.append(data.linear_acceleration.z)
    datalist.append(imudata)
    imu_buf.release()

def imuVisualize(data):
    fig = plt.figure()
    ax = plt.axes()
    i=0
    xlist=[]
    ylist=[]
    zlist=[]
    index=[]
    for item in data:
        index.append(i)
        i+=1
        xlist.append(item[0])
        ylist.append(item[1])
        zlist.append(item[2])
    plt.plot(index, xlist,color='r')
    plt.plot(index, ylist, color='g')
    plt.plot(index, zlist, color='b')

fig = plt.figure()
ax = plt.axes()
def imuVisualize2(data,i):
    xlist=[]
    ylist=[]
    zlist=[]
    index=[]
    for item in data:
        index.append(i)
        i+=1
        xlist.append(item[0])
        ylist.append(item[1])
        zlist.append(item[2])
    plt.plot(index, xlist,color='r')
    plt.plot(index, ylist, color='g')
    plt.plot(index, zlist, color='b')
    return i


def get_imu_record(path):
    list_files = []
    #path_list = os.listdir(path)
    #path_list.sort(key=lambda x:int(x.split()))
    if os.path.exists(path):
        # 遍历文件夹，找到.npy文件
        for root, ds, fs in os.walk(path):
            for f in fs:
                if f.endswith('.npy'):
                    fullname = os.path.join(root, f)
                    list_files.append(fullname)
    else:
        print('this path not exist')
    list_files.sort(key=lambda x:int(x.split('/')[2].split('.')[0]))
    return list_files

if __name__ == "__main__":
    rospy.init_node('IMUhandle',anonymous=True)
    imusub = rospy.Subscriber("/imu0", Imu,handleIMU, queue_size=20)
    i=0
    FirstPublish=True
    while True:
        if FirstPublish:
            if len(datalist) >= 600:
                print("first public")
                # imuPubTask = threading.Thread(target=initialPublishImuByGroup)
                #
                # imuPubTask.start()
                #
                # imuPubTask.join()
                # FirstPublish = False
                imu_buf.acquire()
                datalist=[]
                imu_buf.release()
                i = i + 1
                FirstPublish=False
                print("done with " + str(i))
                print("first public done")
        elif len(datalist)>=300:
            imu_buf.acquire()
            check=[]
            check=datalist
            datalist=[]
            imu_buf.release()
            if len(check)==0:
                break;
            print("check size :"+str(len(check)))
            path = os.path.join("../imu_record", str(i))
            os.makedirs(path)
            file_name = os.path.join(path, str(i) + ".npy")
            np.save(file_name, np.array(check))
            i+=1

            check = []
            print("done with "+str(i))
            if i>=121:
                break

    #imuVisualize(datalist)
    files= get_imu_record("../imu_record")
    i=1
    for file in files:
        # 加载动作数据
        data_train = np.load(file)
        # 计算测试动作与存储动作之间的距离
        imuVisualize(data_train)
        plt.savefig("./result/" + str(i) + ".png")  # ！！！
        i=i+1
        #score = estimate_dtw(np.array(data_train), np.array(datalist), dis_func=dis_ACC)
        #print(score[0])
    print("done")
    #plt.show()
    rospy.spin()
